{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "\n",
    "import sys\n",
    "import os\n",
    "import numpy as np\n",
    "from easydict import EasyDict as edict\n",
    "from glob import glob\n",
    "from tqdm import tqdm\n",
    "\n",
    "sys.path.append('../insightface/src/align')\n",
    "sys.path.append('../insightface/src/common')\n",
    "sys.path.append('../insightface/src/megaface')\n",
    "\n",
    "import detect_face\n",
    "import cv2\n",
    "\n",
    "import mxnet as mx\n",
    "import face_preprocess\n",
    "\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"   # see issue #152\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading /mnt/ssd2/David/faces/insightface/recognition/models/r100-arcface-emore/model 1\n"
     ]
    }
   ],
   "source": [
    "nets = []\n",
    "image_shape = [3,112,112]\n",
    "models_list = ['/mnt/ssd2/David/faces/insightface/recognition/models/r100-arcface-emore/model,1']\n",
    "\n",
    "for model in models_list:\n",
    "    vec = model.split(',')\n",
    "    assert len(vec)>1\n",
    "    prefix = vec[0]\n",
    "    epoch = int(vec[1])\n",
    "    print('loading',prefix, epoch)\n",
    "    net = edict()\n",
    "    net.ctx = mx.gpu(0)\n",
    "    net.sym, net.arg_params, net.aux_params = mx.model.load_checkpoint(prefix, epoch)\n",
    "    \n",
    "    all_layers = net.sym.get_internals()\n",
    "    \n",
    "    net.sym = all_layers['fc1_output']\n",
    "    net.model = mx.mod.Module(symbol=net.sym, context=net.ctx, label_names = None)\n",
    "    net.model.bind(data_shapes=[('data', (1, 3, image_shape[1], image_shape[2]))])\n",
    "    net.model.set_params(net.arg_params, net.aux_params)\n",
    "    nets.append(net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_feature(image_list, net):\n",
    "    images = []\n",
    "    for image_path in image_list:\n",
    "        img = face_preprocess.read_image(image_path, mode='rgb')\n",
    "        img = np.transpose( img, (2,0,1) )\n",
    "        images.append(img)\n",
    "\n",
    "    input_blob = np.array(images)\n",
    "    data = mx.nd.array(input_blob)\n",
    "    db = mx.io.DataBatch(data=(data,))\n",
    "    net.model.forward(db, is_train=False)\n",
    "    embedding = net.model.get_outputs()[0].asnumpy()\n",
    "        \n",
    "    return embedding"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Process images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Images: 5822653\n"
     ]
    }
   ],
   "source": [
    "images_folder = \"/mnt/ssd2/Datasets/faces_emore_images\"\n",
    "dst_folder = \"/mnt/ssd2/Datasets/faces_emore_labels\"\n",
    "images = glob(os.path.join(images_folder, \"*.png\"))\n",
    "np.random.shuffle(images)\n",
    "print(\"Images:\", len(images))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5822653/5822653 [5:36:41<00:00, 288.23it/s]  \n"
     ]
    }
   ],
   "source": [
    "batch_size = 50\n",
    "files = []\n",
    "\n",
    "for filename in tqdm(images):\n",
    "    files.append(filename)\n",
    "        \n",
    "    if len(files) == batch_size:\n",
    "        vectors = get_feature(files, nets[0])\n",
    "        for path, vector in zip(files, vectors):\n",
    "            np.save(path.replace(images_folder, dst_folder)+\"_fc1\", vector)\n",
    "        files = []\n",
    "\n",
    "vectors = get_feature(files, nets[0])\n",
    "for path, vector in zip(files, vectors):\n",
    "    np.save(path.replace(images_folder, dst_folder)+\"_fc1\", vector)\n",
    "files = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"finish\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
